import tkinter as tk
from tkinter import filedialog
from collections import OrderedDict
import numpy as np
import dlib
import cv2

# 全局变量，用于存储选定的模型和图片路径
shape_predictor_path = ""
image_path = ""

def select_model():
    """选择dat模型文件"""
    global shape_predictor_path
    shape_predictor_path = filedialog.askopenfilename(
        title="选择dat模型文件",
        filetypes=[("DAT files", "*.dat"), ("All files", "*")]
    )
    model_label.config(text=shape_predictor_path if shape_predictor_path else "未选择模型")

def select_image():
    """选择输入图片"""
    global image_path
    image_path = filedialog.askopenfilename(
        title="选择图片文件",
        filetypes=[("Image files", "*.jpg *.jpeg *.png *.bmp"), ("All files", "*")]
    )
    image_label.config(text=image_path if image_path else "未选择图片")

def run_detection():
    """执行人脸检测与关键点定位"""
    if not shape_predictor_path or not image_path:
        status_label.config(text="请先选择模型文件和图片文件")
        return

    status_label.config(text="运行中...(按下N查看下一个ROI)")
    root.update()

    # 定义面部关键点区域的索引（68个关键点版）
    FACIAL_LANDMARKS_68_IDXS = OrderedDict([
        ("mouth", (48, 68)),
        ("right_eyebrow", (17, 22)),
        ("left_eyebrow", (22, 27)),
        ("right_eye", (36, 42)),
        ("left_eye", (42, 48)),
        ("nose", (27, 36)),
        ("jaw", (0, 17))
    ])

    def shape_to_np(shape, dtype="int"):
        # 将dlib的关键点对象转换为NumPy数组
        coords = np.zeros((shape.num_parts, 2), dtype=dtype)
        for i in range(0, shape.num_parts):
            coords[i] = (shape.part(i).x, shape.part(i).y)
        return coords

    def visualize_facial_landmarks(image, shape, colors=None, alpha=0.75):
        # 在图像上绘制各区域的关键点和轮廓
        overlay = image.copy()
        output = image.copy()
        if colors is None:
            colors = [(19, 199, 109), (79, 76, 240), (230, 159, 23),
                      (168, 100, 168), (158, 163, 32),
                      (163, 38, 32), (180, 42, 220)]
        for (i, name) in enumerate(FACIAL_LANDMARKS_68_IDXS.keys()):
            (j, k) = FACIAL_LANDMARKS_68_IDXS[name]
            pts = shape[j:k]
            if name == "jaw":
                # 对下颌部分用线条连接
                for l in range(1, len(pts)):
                    ptA = tuple(pts[l - 1])
                    ptB = tuple(pts[l])
                    cv2.line(overlay, ptA, ptB, colors[i], 2)
            else:
                # 其他区域使用凸包填充
                hull = cv2.convexHull(pts)
                cv2.drawContours(overlay, [hull], -1, colors[i], -1)
        cv2.addWeighted(overlay, alpha, output, 1 - alpha, 0, output)
        return output

    # 加载dlib人脸检测器和预训练的关键点预测器
    detector = dlib.get_frontal_face_detector()
    predictor = dlib.shape_predictor(shape_predictor_path)

    # 读取并预处理输入图片
    image = cv2.imread(image_path)
    (h, w) = image.shape[:2]
    width = 500
    r = width / float(w)
    dim = (width, int(h * r))
    image = cv2.resize(image, dim, interpolation=cv2.INTER_AREA)
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 人脸检测
    rects = detector(gray, 1)
    for (i, rect) in enumerate(rects):
        # 获取人脸关键点
        shape = predictor(gray, rect)
        shape = shape_to_np(shape)
        # 遍历每个面部区域进行展示
        for (name, (i_idx, j_idx)) in FACIAL_LANDMARKS_68_IDXS.items():
            clone = image.copy()
            cv2.putText(clone, name, (10, 30), cv2.FONT_HERSHEY_SIMPLEX,
                        0.7, (0, 0, 255), 2)
            for (x, y) in shape[i_idx:j_idx]:
                cv2.circle(clone, (x, y), 3, (0, 0, 255), -1)
            (x, y, w_box, h_box) = cv2.boundingRect(np.array([shape[i_idx:j_idx]]))
            roi = image[y:y + h_box, x:x + w_box]
            (roi_h, roi_w) = roi.shape[:2]
            roi_width = 250
            r_roi = roi_width / float(roi_w)
            dim_roi = (roi_width, int(roi_h * r_roi))
            roi = cv2.resize(roi, dim_roi, interpolation=cv2.INTER_AREA)
            cv2.imshow("ROI", roi)
            cv2.imshow("Image", clone)
            cv2.waitKey(0)
        # 展示所有区域的可视化效果
        output = visualize_facial_landmarks(image, shape)
        cv2.imshow("Image", output)
        cv2.waitKey(0)
    status_label.config(text="检测完成")

# 构建GUI界面
root = tk.Tk()
root.title("面部关键点检测GUI")

# 设置窗口尺寸并使窗口居中显示
window_width = 600
window_height = 400
screen_width = root.winfo_screenwidth()
screen_height = root.winfo_screenheight()
x_cordinate = int((screen_width / 2) - (window_width / 2))
y_cordinate = int((screen_height / 2) - (window_height / 2))
root.geometry(f"{window_width}x{window_height}+{x_cordinate}+{y_cordinate}")

# 设置较大的字体和按钮尺寸
button_font = ("Arial", 14)
label_font = ("Arial", 12)

select_model_btn = tk.Button(root, text="选择dat模型", command=select_model, font=button_font, width=20)
select_model_btn.pack(pady=10)

model_label = tk.Label(root, text="未选择模型", font=label_font)
model_label.pack(pady=5)

select_image_btn = tk.Button(root, text="选择图片", command=select_image, font=button_font, width=20)
select_image_btn.pack(pady=10)

image_label = tk.Label(root, text="未选择图片", font=label_font)
image_label.pack(pady=5)

run_btn = tk.Button(root, text="运行检测", command=run_detection, font=button_font, width=20)
run_btn.pack(pady=20)

status_label = tk.Label(root, text="等待操作", font=label_font)
status_label.pack(pady=5)

root.mainloop()